import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import fpgrowth, association_rules
import matplotlib.pyplot as plt
from graphviz import Digraph
import pygraphviz as pgv
from mlxtend.frequent_patterns import fpgrowth
from mlxtend.frequent_patterns import association_rules
from mlxtend.frequent_patterns import fpmax
from mlxtend.frequent_patterns import fpgrowth
from collections import defaultdict
import networkx as nx

# 构建FP树
def build_fp_tree(patterns):

    sorted_patterns = {k: v for k, v in sorted(patterns.items(), key=lambda item: item[1], reverse=True)}
    tree = defaultdict(dict)
    for transaction, count in sorted_patterns.items():
        current = tree
        for item in transaction:
            if item not in current:
                current[item] = {'count': count, 'children': {}}
            else:
                current[item]['count'] += count
            current = current[item]['children']
    return tree

# 可视化FP树
def visualize_fp_tree(tree,rules, output_file='fp_tree_visualization.png'):
    G = pgv.AGraph(strict=False, directed=True)
    G.graph_attr['rankdir'] = 'TB'
    G.add_node('null', label='', style='invisible')
    build_graph_with_rules(G, 'null', tree,rules)
    G.layout(prog='dot')
    G.draw(output_file,format='svg')

def build_graph_with_rules(G, parent, tree,rules):
    """根据关联规则绘制FP树"""
    for item, values in tree.items():
        
        node_id = f"{parent}_{item}"
        G.add_node(node_id, label=f"{item} ({values['count']})")
        for index,row in rules.iterrows():

            if set(row['antecedents']).issubset(set(node_id.split('_')[1:])):
                right_item=row['consequents'] if len(row['consequents'])==1 else tuple(row['consequents'])
                right_node_id=f"{node_id}_{right_item}"
                G.add_node(right_node_id,label=f"({values['count']})",style='filled',fillcolor='lightblue')
                G.add_edge(node_id,right_node_id,label=f"conf={row['confidence']:.2f}")
        G.add_edge(parent, node_id)
        build_graph_with_rules(G, node_id, values['children'],row)

def write_results_to_file(rules):
    """将结果写入到文件中"""

    # 写入关联规则
    with open('fptree_association_rules.txt', 'w') as f:
        f.write("Association Rules\n")
        for index,row in rules.iterrows():
            f.write(f"{row['antecedents']} --> {row['consequents']}, Confidence: {row['confidence']:.2f}, support:{row['support']}\n")
    print("已将结果保存到文件")



# 主程序
# 读取CSV文件
df = pd.read_csv("D:\Lenovo\Desktop\云南大学\空间数据挖掘\实验数据\实验数据5.csv", header=None)
pre_transactions = df.values.tolist()
transactions=[[x for x in subset  if pd.isnull(x)==False] for subset in pre_transactions]


# 使用TransactionEncoder将交易记录转换为适合FP-Growth算法的格式
te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df = pd.DataFrame(te_ary, columns=te.columns_)

# 使用FP-Growth算法生成频繁项集
frequent_itemsets = fpgrowth(df, min_support=0.2, use_colnames=True)

# 生成关联规则，生成一个DataFrame，其中每一行都是一个Series，可以理解为每一行都是一个允许值重复的字典
rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.7)

# 打印关联规则及其置信度
print(rules[['antecedents', 'consequents', 'support', 'confidence', 'lift']])
rules=rules.sort_values(by='confidence',ascending=False,axis=0)
write_results_to_file(rules)


# # 定义FP-tree节点类
# class TreeNode:
#     def __init__(self, item, count, parent):
#         self.item = item
#         self.count = count
#         self.parent = parent
#         self.children = {}
#         self.nodeLink = None

#     def increment(self, count):
#         self.count += count


# # 构建FP-tree
# def construct_tree(transactions, min_support):
#     header_table = {}
#     for transaction in transactions:
#         for item in transaction:
#             header_table[item] = header_table.get(item, 0) + 1

#     header_table = {k: v for k, v in header_table.items() if v >= min_support}
#     freq_item_set = set(header_table.keys())

#     if len(freq_item_set) == 0:
#         return None, None

#     for k in header_table:
#         header_table[k] = [header_table[k], None]

#     tree_root = TreeNode('Null', 1, None)

#     for transaction, count in transactions.items():
#         local_d = {}
#         for item in transaction:
#             if item in freq_item_set:
#                 local_d[item] = header_table[item][0]
#         if len(local_d) > 0:
#             ordered_items = [v[0] for v in sorted(local_d.items(), key=lambda p: p[1], reverse=True)]
#             update_tree(ordered_items, tree_root, header_table, count)

#     return tree_root, header_table

# def update_tree(items, node, header_table, count):
#     if items[0] in node.children:
#         node.children[items[0]].increment(count)
#     else:
#         node.children[items[0]] = TreeNode(items[0], count, node)
#         if header_table[items[0]][1] is None:
#             header_table[items[0]][1] = node.children[items[0]]
#         else:
#             update_header(header_table[items[0]][1], node.children[items[0]])

#     if len(items) > 1:
#         update_tree(items[1:], node.children[items[0]], header_table, count)

# def update_header(node, target_node):
#     while node.nodeLink is not None:
#         node = node.nodeLink
#     node.nodeLink = target_node

# # FP-tree的可视化
# def visualize_fp_tree(tree_root):
#     dot = Digraph()

#     def add_nodes_edges(node, parent_name):
#         for child in node.children.values():
#             node_name = f'{child.item} ({child.count})'
#             dot.node(node_name)
#             dot.edge(parent_name, node_name)
#             add_nodes_edges(child, node_name)
    
#     root_name = f'{tree_root.item} ({tree_root.count})'
#     dot.node(root_name)
#     dot.attr(rankdir='TB')
#     add_nodes_edges(tree_root, root_name)
#     dot.render('fp_tree', format='png')

# # 将关联规则转为transaction_dict
# transaction_dict = {}
# for _, row in rules.iterrows():
#     transaction = tuple(row['antecedents']) + tuple(row['consequents'])
#     if transaction in transaction_dict:
#         transaction_dict[transaction] += 1
#     else:
#         transaction_dict[transaction] = 1




# # 构建并可视化FP-tree
# fp_tree=build_fp_tree(transaction_dict)
# visualize_fp_tree(fp_tree,rules,output_file='fp树可视化.svg')
# tree_root, header_table = construct_tree(transaction_dict, min_support)
# if tree_root:
#     visualize_fp_tree(tree_root)
# else:
#     print("No frequent itemsets found.")
